#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
简化版Milvus查询示例
"""
import sys
import os
from pathlib import Path

# 添加项目根目录到系统路径
project_root = Path(__file__).parent.parent.parent.parent
sys.path.append(str(project_root))

try:
    from pymilvus import MilvusClient, Collection, connections, utility
    print("成功导入pymilvus库")
except ImportError as e:
    print(f"错误: 导入pymilvus失败 - {e}")
    sys.exit(1)

# 添加backend目录到系统路径（仅用于测试脚本）
backend_dir = Path(__file__).parent.parent.parent
sys.path.insert(0, str(backend_dir))

try:
    from app.utils.config_manager import config
    milvus_uri = config.milvus_config.get("uri")
    milvus_token = config.milvus_config.get("token")
    collection_name = config.milvus_config.get("collection_name", "node_vectors")
    print(f"读取配置: URI={milvus_uri}, 集合={collection_name}")
except ImportError as e:
    print(f"错误: 导入配置管理器失败 - {e}")
    # 使用默认值
    milvus_uri = "http://localhost:19530"
    milvus_token = ""
    collection_name = "node_vectors"

def test_collection():
    """测试集合信息查询"""
    print("\n=== 测试MilvusClient类 ===")
    
    client = None
    try:
        print(f"尝试连接: {milvus_uri}")
        client = MilvusClient(uri=milvus_uri, token=milvus_token)
        print("MilvusClient连接成功")
        
        # 检查客户端对象的可用方法
        print("\nMilvusClient可用方法:")
        for method_name in [m for m in dir(client) if not m.startswith('_')]:
            print(f"  - {method_name}")
        
        # 获取集合列表
        collections = client.list_collections()
        print(f"\n发现集合 ({len(collections)}):", collections)
        
        if collection_name not in collections:
            print(f"错误: 集合 {collection_name} 不存在")
            return
        
        # 获取集合信息
        try:
            collection_info = client.describe_collection(collection_name)
            print(f"\n集合信息: {collection_name}")
            for key, value in collection_info.items():
                if isinstance(value, dict) or isinstance(value, list):
                    print(f"  - {key}: [复杂数据]")
                else:
                    print(f"  - {key}: {value}")
            
            # 尝试提取字段信息
            fields = []
            if "schema" in collection_info and "fields" in collection_info["schema"]:
                fields = collection_info["schema"]["fields"]
            elif "fields" in collection_info:
                fields = collection_info["fields"]
            
            if fields:
                print("\n字段列表:")
                for field in fields:
                    print(f"  - {field.get('name', 'unknown')}: {field.get('type', 'unknown')}")
                    
            # 尝试查询一条记录
            try:
                print(f"\n尝试从 {collection_name} 查询一条数据:")
                data = client.query(
                    collection_name=collection_name,
                    filter="",
                    output_fields=[field.get('name') for field in fields if field.get('name') != "vector"],
                    limit=1
                )
                print(f"查询结果: {data}")
            except Exception as e:
                print(f"查询示例记录失败: {e}")
                
        except Exception as e:
            print(f"获取集合信息失败: {e}")
            
    except Exception as e:
        print(f"使用MilvusClient时出错: {e}")
    
    print("\n=== 测试传统API ===")
    try:
        connections.connect("default", uri=milvus_uri, token=milvus_token)
        print("传统连接成功")
        
        col_list = utility.list_collections()
        print(f"集合列表: {col_list}")
        
        if collection_name in col_list:
            collection = Collection(collection_name)
            print(f"集合加载成功: {collection_name}")
            print(f"实体数量: {collection.num_entities}")
            
            schema = collection.schema
            print("\n字段列表(传统API):")
            for field in schema.fields:
                print(f"  - {field.name}: {field.dtype}")
    except Exception as e:
        print(f"使用传统API时出错: {e}")
            
if __name__ == "__main__":
    test_collection() 